use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use prost::Message;
use futures::future::BoxFuture;
use crate::vector_proto::{CollectionProto, VectorProto};


pub trait StorageTrait {
    fn save(&self, vectors: &HashMap<String, Vector>) -> BoxFuture<'_, std::io::Result<()>>;
    fn load(&self) -> BoxFuture<'_, std::io::Result<HashMap<String, Vector>>>;
}

pub struct JsonStorage {
    storage_path: String,
}

impl JsonStorage {
    pub fn new(storage_path: &str) -> Self {
        JsonStorage {
            storage_path: storage_path.to_string(),
        }
    }
}

impl StorageTrait for JsonStorage {
    fn save(&self, vectors: &HashMap<String, Vector>) -> BoxFuture<'_, std::io::Result<()>> {
        let path = self.storage_path.clone();
        let vectors = vectors.clone();
        Box::pin(async move {
            let serialized = serde_json::to_string(&vectors)?;
            tokio::fs::write(&path, serialized).await
        })
    }

    fn load(&self) -> BoxFuture<'_, std::io::Result<HashMap<String, Vector>>> {
        let path = self.storage_path.clone();
        Box::pin(async move {
            let data = tokio::fs::read_to_string(&path).await?;
            Ok(serde_json::from_str(&data)?)
        })
    }
}

pub struct ProtoStorage {
    storage_path: String,
}

impl ProtoStorage {
    pub fn new(storage_path: &str) -> Self {
        ProtoStorage {
            storage_path: storage_path.to_string(),
        }
    }
}

impl StorageTrait for ProtoStorage {
    fn save(&self, vectors: &HashMap<String, Vector>) -> BoxFuture<'_, std::io::Result<()>> {
        let path = self.storage_path.clone();
        let vectors = vectors.clone();
        Box::pin(async move {
            let mut buf = Vec::new();
            let mut collection_proto = CollectionProto::default();
            
            for (id, vector) in vectors.iter() {
                let mut vector_proto = VectorProto::default();
                vector_proto.set_id(id.clone());
                vector_proto.set_embedding(vector.embedding.clone());
                vector_proto.set_metadata(serde_json::to_string(&vector.metadata).unwrap());
                collection_proto.mut_vectors().push(vector_proto);
            }
            
            collection_proto.write_to_vec(&mut buf)?;
            tokio::fs::write(&path, &buf).await
        })
    }

    fn load(&self) -> BoxFuture<'_, std::io::Result<HashMap<String, Vector>>> {
        let path = self.storage_path.clone();
        Box::pin(async move {
            let data = tokio::fs::read(&path).await?;
            let mut collection_proto = CollectionProto::decode(&data[..])?;
            
            let mut vectors = HashMap::new();
            for vector_proto in collection_proto.take_vectors().into_iter() {
                let metadata: serde_json::Value = serde_json::from_str(&vector_proto.get_metadata())
                    .unwrap_or(serde_json::Value::Null);
                
                vectors.insert(
                    vector_proto.get_id().to_string(),
                    Vector {
                        id: vector_proto.get_id().to_string(),
                        embedding: vector_proto.get_embedding().to_vec(),
                        metadata,
                    },
                );
            }
            
            Ok(vectors)
        })
    }
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Vector {
    pub id: String,
    pub embedding: Vec<f32>,
    pub metadata: serde_json::Value, // 用户自定义字段
}

#[derive(Clone)]
pub struct Collection {
    name: String,
    vectors: Arc<RwLock<HashMap<String, Vector>>>, // 内存中的数据
    storage: Arc<dyn StorageTrait + Send + Sync>, // 存储实现
}

impl Collection {
    pub fn new_with_json_storage(name: &str, storage_path: &str) -> Self {
        log::info!("{}", storage_path);
        Collection {
            name: name.to_string(),
            vectors: Arc::new(RwLock::new(HashMap::new())),
            storage: Arc::new(JsonStorage::new(storage_path)),
        }
    }
    
    pub fn new_with_proto_storage(name: &str, storage_path: &str) -> Self {
        log::info!("{}", storage_path);
        Collection {
            name: name.to_string(),
            vectors: Arc::new(RwLock::new(HashMap::new())),
            storage: Arc::new(ProtoStorage::new(storage_path)),
        }
    }

    pub fn get_name(&self) -> &str {
        &self.name
    }

    pub async fn insert(&self, vector: Vector) {
        {
            // Acquire the write lock and insert the vector
            let mut vectors = self.vectors.write().await;
            vectors.insert(vector.id.clone(), vector);
        } // Release the write lock here

        // Save to disk outside the lock to avoid deadlocks
        if let Err(e) = self.save_to_disk().await {
            log::error!("Failed to save collection to disk: {}", e);
        } else {
            log::info!("Collection saved to disk successfully.");
        }
    }

    /// Bulk insert multiple vectors into the collection.
    pub async fn bulk_insert(&self, vectors: Vec<Vector>) {
        {
            // Acquire the write lock and insert all vectors
            let mut collection_vectors = self.vectors.write().await;
            for vector in vectors {
                collection_vectors.insert(vector.id.clone(), vector);
            }
        } // Release the write lock here

        // Save to disk outside the lock to avoid deadlocks
        if let Err(e) = self.save_to_disk().await {
            log::error!("Failed to save collection to disk: {}", e);
        } else {
            log::info!("Collection saved to disk successfully after bulk insert.");
        }
    }

    pub async fn get(&self, id: &str) -> Option<Vector> {
        let vectors = self.vectors.read().await;
        vectors.get(id).cloned()
    }

    pub async fn delete(&self, id: &str) -> Option<Vector> {
        let result = {
            // Acquire the write lock and remove the vector
            let mut vectors = self.vectors.write().await;
            vectors.remove(id)
        }; // Release the write lock here

        // Save to disk outside the lock to avoid deadlocks
        if let Err(e) = self.save_to_disk().await {
            log::error!("Failed to save collection to disk: {}", e);
        } else {
            log::info!("Collection saved to disk successfully.");
        }

        result
    }

    /// Bulk delete multiple vectors from the collection.
    pub async fn bulk_delete(&self, ids: Vec<String>) -> Vec<Option<Vector>> {
        let mut results = Vec::new();
        {
            // Acquire the write lock and remove all specified vectors
            let mut vectors = self.vectors.write().await;
            for id in ids {
                results.push(vectors.remove(&id));
            }
        } // Release the write lock here

        // Save to disk outside the lock to avoid deadlocks
        if let Err(e) = self.save_to_disk().await {
            log::error!("Failed to save collection to disk after bulk delete: {}", e);
        } else {
            log::info!("Collection saved to disk successfully after bulk delete.");
        }

        results
    }

    pub async fn save_to_disk(&self) -> std::io::Result<()> {
        let vectors = self.vectors.read().await;
        self.storage.save(&vectors).await
    }

    pub async fn load_from_disk(&self) -> std::io::Result<()> {
        let vectors = self.storage.load().await?;
        let mut collection_vectors = self.vectors.write().await;
        *collection_vectors = vectors;
        Ok(())
    }

    /// Search for the top `k` most similar vectors based on cosine similarity.
    pub async fn search(&self, query: &[f32], k: usize) -> Vec<(String, f32)> {
        let vectors = self.vectors.read().await;

        // Compute cosine similarity for all vectors
        let mut similarities: Vec<(String, f32)> = vectors
            .iter()
            .map(|(id, vector)| {
                let similarity = cosine_similarity(query, &vector.embedding);
                (id.clone(), similarity)
            })
            .collect();

        // Sort by similarity in descending order
        similarities.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        // Return the top `k` results
        similarities.into_iter().take(k).collect()
    }
}

pub struct VectorDB {
    collections: Arc<RwLock<HashMap<String, Collection>>>,
    storage_type: String,
}

impl VectorDB {
    pub fn new(storage_type: String) -> Self {
        VectorDB {
            collections: Arc::new(RwLock::new(HashMap::new())),
            storage_type,
        }
    }

    /// Initialize the database by loading collections from disk.
    pub async fn initialize(&self, storage_dir: &str) -> std::io::Result<()> {
        let mut paths = tokio::fs::read_dir(storage_dir).await?;
        let mut collections = self.collections.write().await;

        while let Some(entry) = paths.next_entry().await? {
            let path = entry.path();
            if let Some(ext) = path.extension().and_then(|ext| ext.to_str()) {
                if ext == "json" || ext == "proto" || ext == "bin" || ext == "dat" {
                    if let Some(name) = path.file_stem().and_then(|stem| stem.to_str()) {
                        let collection = if ext == "json" {
                            Collection::new_with_json_storage(name, path.to_str().unwrap())
                        } else {
                            Collection::new_with_proto_storage(name, path.to_str().unwrap())
                        };

                        match collection.load_from_disk().await {
                            Ok(_) => {
                                log::info!("Loaded collection '{}'", name);
                                collections.insert(name.to_string(), collection);
                            }
                            Err(e) => {
                                log::warn!(
                                    "Failed to load collection '{}': {}. Skipping...",
                                    name,
                                    e
                                );
                            }
                        }
                    }
                }
            }
        }

        log::info!("VectorDB initialized with collections from '{}'", storage_dir);
        Ok(())
    }

    pub async fn create_collection(&self, name: &str) {
        let (storage_path, use_proto) = match self.storage_type.as_str() {
            "proto" => (format!("./collections/{}.proto.dat", name), true),
            _ => (format!("./collections/{}.json", name), false)
        };

        let collection = if use_proto {
            Collection::new_with_proto_storage(name, &storage_path)
        } else {
            Collection::new_with_json_storage(name, &storage_path)
        };
        
        let mut collections = self.collections.write().await;
        collections.insert(name.to_string(), collection);
    }

    pub async fn get_collection(&self, name: &str) -> Option<Collection> {
        let collections = self.collections.read().await;
        collections.get(name).cloned()
    }

    /// Get a collection by name, creating it if necessary.
    pub async fn get_or_load_collection(&self, name: &str) -> Result<Collection, String> {
        let mut collections = self.collections.write().await;

        // Check if the collection is already in memory
        if let Some(collection) = collections.get(name) {
            return Ok(collection.clone());
        }

        // Attempt to load the collection from disk
        let proto_path = format!("./collections/{}.proto.dat", name);
        let json_path = format!("./collections/{}.json", name);

        let collection = if std::path::Path::new(&proto_path).exists() {
            let collection = Collection::new_with_proto_storage(name, &proto_path);
            if let Err(e) = collection.load_from_disk().await {
                log::error!("Failed to load collection '{}': {}", name, e);
                return Err(format!("Failed to load collection '{}': {}", name, e));
            }
            log::info!("Loaded collection '{}' from disk (proto format)", name);
            collection
        } else if std::path::Path::new(&json_path).exists() {
            let collection = Collection::new_with_json_storage(name, &json_path);
            if let Err(e) = collection.load_from_disk().await {
                log::error!("Failed to load collection '{}': {}", name, e);
                return Err(format!("Failed to load collection '{}': {}", name, e));
            }
            log::info!("Loaded collection '{}' from disk (json format)", name);
            collection
        } else {
            // Create a new collection if it does not exist
            log::info!("Collection '{}' not found. Creating a new one.", name);
            let storage_path = match self.storage_type.as_str() {
                "proto" => format!("./collections/{}.proto.dat", name),
                _ => format!("./collections/{}.json", name),
            };
            if self.storage_type == "proto" {
                Collection::new_with_proto_storage(name, &storage_path)
            } else {
                Collection::new_with_json_storage(name, &storage_path)
            }
        };

        collections.insert(name.to_string(), collection.clone());
        Ok(collection)
    }

    pub async fn insert(&self, collection_name: String, vector: Vector) -> Result<(), String> {
        let collection = self.get_or_load_collection(&collection_name).await?;
        log::info!("Inserting vector into collection '{}'", collection_name);
        collection.insert(vector).await;
        log::info!("Vector inserted into collection '{}'", collection_name);
        Ok(())
    }

    /// Bulk insert vectors into a specific collection.
    pub async fn bulk_insert(
        &self,
        collection_name: String,
        vectors: Vec<Vector>,
    ) -> Result<(), String> {
        let collection = self.get_or_load_collection(&collection_name).await?;
        log::info!(
            "Bulk inserting {} vectors into collection '{}'",
            vectors.len(),
            collection_name
        );
        collection.bulk_insert(vectors).await;
        log::info!("Bulk insert completed for collection '{}'", collection_name);
        Ok(())
    }

    pub async fn delete(&self, collection_name: String, id: String) -> Result<Option<Vector>, String> {
        let collection = self.get_or_load_collection(&collection_name).await?;
        log::info!("Deleting vector '{}' from collection '{}'", id, collection_name);
        Ok(collection.delete(&id).await)
    }

    /// Bulk delete vectors from a specific collection.
    pub async fn bulk_delete(&self, collection_name: String, ids: Vec<String>) -> Result<(), String> {
        let collection = self.get_or_load_collection(&collection_name).await?;
        log::info!(
            "Bulk deleting {} vectors from collection '{}'",
            ids.len(),
            collection_name
        );
        collection.bulk_delete(ids).await;
        log::info!("Bulk delete completed for collection '{}'", collection_name);
        Ok(())
    }

    pub async fn search(
        &self,
        collection_name: String,
        query: Vec<f32>,
        k: usize,
    ) -> Result<Vec<(String, f32)>, String> {
        let collection = self.get_or_load_collection(&collection_name).await?;
        log::info!(
            "Searching in collection '{}' with query vector of length {}",
            collection_name,
            query.len()
        );
        Ok(collection.search(&query, k).await)
    }

    /// Close a collection by saving its data to disk and removing it from memory.
    pub async fn close_collection(&self, collection_name: &str) -> Result<(), String> {
        let mut collections = self.collections.write().await;
        if let Some(collection) = collections.remove(collection_name) {
            log::info!("Closing collection '{}'", collection_name);
            if let Err(e) = collection.save_to_disk().await {
                log::error!("Failed to save collection '{}': {}", collection_name, e);
                return Err(format!("Failed to save collection '{}': {}", collection_name, e));
            }
            log::info!("Collection '{}' closed and removed from memory", collection_name);
            Ok(())
        } else {
            Err(format!("Collection '{}' not found", collection_name))
        }
    }
}

fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
    let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();

    if norm_a == 0.0 || norm_b == 0.0 {
        0.0 // Return 0 similarity if either vector has zero magnitude
    } else {
        dot_product / (norm_a * norm_b)
    }
}